package com.yl.request.common.glaencrypt.sm2;

import org.bouncycastle.asn1.*;
import org.bouncycastle.crypto.*;
import org.bouncycastle.crypto.digests.SM3Digest;
import org.bouncycastle.crypto.params.*;
import org.bouncycastle.math.ec.*;
import org.bouncycastle.util.*;
import java.io.ByteArrayOutputStream;
import java.math.BigInteger;
import java.security.SecureRandom;


public class SM2Engine {
    private final Digest digest;
    private final Mode mode;
    private boolean forEncryption;
    private ECKeyParameters ecKey;
    private ECDomainParameters ecParams;
    private int curveLength;
    private SecureRandom random;

    public SM2Engine() {
        this(new SM3Digest());
    }

    public SM2Engine(Mode mode) {
        this(new SM3Digest(), mode);
    }

    public SM2Engine(Digest digest) {
        this(digest, Mode.C1C2C3);
    }

    public SM2Engine(Digest digest, Mode mode) {
        if (mode == null) {
            throw new IllegalArgumentException("mode cannot be NULL");
        } else {
            this.digest = digest;
            this.mode = mode;
        }
    }

    public void init(boolean forEncryption, CipherParameters param) {
        this.forEncryption = forEncryption;
        if (forEncryption) {
            ParametersWithRandom rParam = (ParametersWithRandom)param;
            this.ecKey = (ECKeyParameters)rParam.getParameters();
            this.ecParams = this.ecKey.getParameters();
            ECPoint s = ((ECPublicKeyParameters)this.ecKey).getQ().multiply(this.ecParams.getH());
            if (s.isInfinity()) {
                throw new IllegalArgumentException("invalid key: [h]Q at infinity");
            }

            this.random = rParam.getRandom();
        } else {
            this.ecKey = (ECKeyParameters)param;
            this.ecParams = this.ecKey.getParameters();
        }

        this.curveLength = (this.ecParams.getCurve().getFieldSize() + 7) / 8;
    }

    public byte[] processBlock(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        return this.forEncryption ? this.encrypt(in, inOff, inLen) : this.decrypt(in, inOff, inLen);
    }

    public int getOutputSize(int inputLen) {
        return 1 + 2 * this.curveLength + inputLen + this.digest.getDigestSize();
    }

    protected ECMultiplier createBasePointMultiplier() {
        return new FixedPointCombMultiplier();
    }

    private byte[] encrypt(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        byte[] c2 = new byte[inLen];
        System.arraycopy(in, inOff, c2, 0, c2.length);
        ECMultiplier multiplier = this.createBasePointMultiplier();

        byte[] c1;
        ECPoint kPB;
        do {
            BigInteger k = this.nextK();
            ECPoint c1P = multiplier.multiply(this.ecParams.getG(), k).normalize();
            c1 = c1P.getEncoded(false);
            kPB = ((ECPublicKeyParameters)this.ecKey).getQ().multiply(k).normalize();
            this.kdf(this.digest, kPB, c2);
        } while(this.notEncrypted(c2, in, inOff));

        byte[] c3 = new byte[this.digest.getDigestSize()];
        this.addFieldElement(this.digest, kPB.getAffineXCoord());
        this.digest.update(in, inOff, inLen);
        this.addFieldElement(this.digest, kPB.getAffineYCoord());
        this.digest.doFinal(c3, 0);
        switch(this.mode) {
            case C1C3C2:
                return Arrays.concatenate(c1, c3, c2);
            default:
                return Arrays.concatenate(c1, c2, c3);
        }
    }

    private byte[] decrypt(byte[] in, int inOff, int inLen) throws InvalidCipherTextException {
        byte[] c1 = new byte[this.curveLength * 2 + 1];
        System.arraycopy(in, inOff, c1, 0, c1.length);
        ECPoint c1P = this.ecParams.getCurve().decodePoint(c1);
        ECPoint s = c1P.multiply(this.ecParams.getH());
        if (s.isInfinity()) {
            throw new InvalidCipherTextException("[h]C1 at infinity");
        } else {
            c1P = c1P.multiply(((ECPrivateKeyParameters)this.ecKey).getD()).normalize();
            int digestSize = this.digest.getDigestSize();
            byte[] c2 = new byte[inLen - c1.length - digestSize];
            if (this.mode == Mode.C1C3C2) {
                System.arraycopy(in, inOff + c1.length + digestSize, c2, 0, c2.length);
            } else {
                System.arraycopy(in, inOff + c1.length, c2, 0, c2.length);
            }

            this.kdf(this.digest, c1P, c2);
            byte[] c3 = new byte[this.digest.getDigestSize()];
            this.addFieldElement(this.digest, c1P.getAffineXCoord());
            this.digest.update(c2, 0, c2.length);
            this.addFieldElement(this.digest, c1P.getAffineYCoord());
            this.digest.doFinal(c3, 0);
            int check = 0;
            int i;
            if (this.mode == Mode.C1C3C2) {
                for(i = 0; i != c3.length; ++i) {
                    check |= c3[i] ^ in[inOff + c1.length + i];
                }
            } else {
                for(i = 0; i != c3.length; ++i) {
                    check |= c3[i] ^ in[inOff + c1.length + c2.length + i];
                }
            }

            Arrays.fill(c1, (byte)0);
            Arrays.fill(c3, (byte)0);
            if (check != 0) {
                Arrays.fill(c2, (byte)0);
                throw new InvalidCipherTextException("invalid cipher text");
            } else {
                return c2;
            }
        }
    }

    private boolean notEncrypted(byte[] encData, byte[] in, int inOff) {
        for(int i = 0; i != encData.length; ++i) {
            if (encData[i] != in[inOff + i]) {
                return false;
            }
        }

        return true;
    }

    private void kdf(Digest digest, ECPoint c1, byte[] encData) {
        int digestSize = digest.getDigestSize();
        byte[] buf = new byte[Math.max(4, digestSize)];
        int off = 0;
        Memoable memo = null;
        Memoable copy = null;
        if (digest instanceof Memoable) {
            this.addFieldElement(digest, c1.getAffineXCoord());
            this.addFieldElement(digest, c1.getAffineYCoord());
            memo = (Memoable)digest;
            copy = memo.copy();
        }

        int xorLen;
        for(int ct = 0; off < encData.length; off += xorLen) {
            if (memo != null) {
                memo.reset(copy);
            } else {
                this.addFieldElement(digest, c1.getAffineXCoord());
                this.addFieldElement(digest, c1.getAffineYCoord());
            }

            ++ct;
            Pack.intToBigEndian(ct, buf, 0);
            digest.update(buf, 0, 4);
            digest.doFinal(buf, 0);
            xorLen = Math.min(digestSize, encData.length - off);
            this.xor(encData, buf, off, xorLen);
        }

    }

    private void xor(byte[] data, byte[] kdfOut, int dOff, int dRemaining) {
        for(int i = 0; i != dRemaining; ++i) {
            data[dOff + i] ^= kdfOut[i];
        }

    }

    private BigInteger nextK() {
        int qBitLength = this.ecParams.getN().bitLength();

        BigInteger k;
        do {
            do {
                k = BigIntegers.createRandomBigInteger(qBitLength, this.random);
            } while(k.equals(BigIntegers.ZERO));
        } while(k.compareTo(this.ecParams.getN()) >= 0);

        return k;
    }

    private void addFieldElement(Digest digest, ECFieldElement v) {
        byte[] p = BigIntegers.asUnsignedByteArray(this.curveLength, v.toBigInteger());
        digest.update(p, 0, p.length);
    }

    public byte[] asn1Toc1c3c2(byte[] asn1tData) throws Exception {
        ASN1InputStream aIn = new ASN1InputStream(asn1tData);
        ASN1Sequence seq = (ASN1Sequence)aIn.readObject();
        BigInteger x = ASN1Integer.getInstance(seq.getObjectAt(0)).getValue();
        BigInteger y = ASN1Integer.getInstance(seq.getObjectAt(1)).getValue();
        byte[] c3 = ASN1OctetString.getInstance(seq.getObjectAt(2)).getOctets();
        byte[] c2 = ASN1OctetString.getInstance(seq.getObjectAt(3)).getOctets();
        ECPoint p = this.ecParams.getCurve().validatePoint(x, y);
        byte[] c1b = p.getEncoded(false);
        ByteArrayOutputStream os = new ByteArrayOutputStream();
        os.write(c1b);
        os.write(c3);
        os.write(c2);
        return os.toByteArray();
    }

    public enum Mode {
        C1C2C3,
        C1C3C2;

        Mode() {
        }
    }
}
